import os
import sys
import logging
import time
import numpy as np
from onnxruntime import InferenceSession
from transformers import T5Tokenizer

from .. import helper
from ..onnx_base import ONNXBase


class ONNXT5EncInfer(ONNXBase):
    def __init__(self, *args) -> None:
        super(ONNXT5EncInfer, self).__init__(*args)
        logging.debug(f"ONNXT5EncInfer init.")
        self.seq_len = 128
        self.tokenizer = T5Tokenizer.from_pretrained(self.model_dir)

        enc_onnx_path = f"{self.model_dir}/encoder.onnx"
        self.enc_sess = InferenceSession(enc_onnx_path, providers=self.providers)
        self.enc_inputs_name = self._get_onnx_input_name(self.enc_sess)
        self.enc_outputs_name = self._get_onnx_output_name(self.enc_sess)

    def __del__(self):
        super(ONNXT5EncInfer, self).__del__()
        logging.debug(f"ONNXT5EncInfer destruct.")

    def _preproc(self, data):
        tokens = self.tokenizer(data, padding="max_length", max_length=self.seq_len, return_tensors="np")
        input_ids = tokens.input_ids
        attention_mask = tokens.attention_mask
        input_ids = input_ids.astype(np.int32)
        attention_mask = attention_mask.astype(np.int32)

        logging.info(f"ids shape:{input_ids.shape} {input_ids.dtype}")
        logging.info(f"mask shape:{attention_mask.shape} {attention_mask.dtype}")

        return input_ids, attention_mask

    def _onnx_infer(self, enc_ids, enc_mask):
        enc_start = time.time()
        enc_hidden_state, = self.enc_sess.run(self.enc_outputs_name, 
                                            {self.enc_inputs_name[0]:enc_ids, 
                                             self.enc_inputs_name[1]:enc_mask})
        enc_time = (time.time() - enc_start) * 1000
        logging.debug(f"shape encoder:{enc_hidden_state.shape}")
        logging.info(f"encoder time:{enc_time:.3f} (ms)")

        return enc_hidden_state

    def enc(self, contents):
        input_ids, attention_mask = self._preproc(contents)
        enc_hidden = self._onnx_infer(input_ids, attention_mask)

        return attention_mask, enc_hidden
